bl_info = {
    "name": "FFD Tools for Blender",
    "author": "You & Assistant",
    "version": (1, 2, 0),
    "blender": (4, 5, 0),
    "location": "View3D > Sidebar > FFD Tools",
    "description": "3ds Max风格的FFD工具，分离物体模式和编辑模式操作",
    "category": "Mesh",
    "support": "COMMUNITY"
}

import bpy
import bmesh
from mathutils import Vector
from bpy.types import Operator, Panel
from bpy.props import EnumProperty, StringProperty

# 确保晶格的scale不为0
min_size = 0.001
def get_bounds_center(obj):
    # 确保包围盒是最新
    bpy.context.view_layer.update()
    
    # 获取局部包围盒的8个点并转换到世界空间
    world_coords = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    
    # 分离各轴坐标
    xs = [v.x for v in world_coords]
    ys = [v.y for v in world_coords]
    zs = [v.z for v in world_coords]
    
    # 计算最小/最大
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    min_z, max_z = min(zs), max(zs)
    
    # 中心 = (min + max) / 2
    center = Vector(((min_x + max_x) / 2,
                     (min_y + max_y) / 2,
                     (min_z + max_z) / 2))
    
    # 尺寸 = max - min（各轴长度，非负）
    size = Vector((max_x - min_x,
                   max_y - min_y,
                   max_z - min_z))
    return size, center
# ==================== 物体模式FFD工具 ====================

class FFD_OT_add_object_ffd(Operator):
    """为整个物体添加FFD"""
    bl_idname = "object.ffd_add_object"
    bl_label = "物体FFD"
    bl_description = "为整个物体添加FFD修改器"
    bl_options = {'REGISTER', 'UNDO'}
    
    resolution: EnumProperty(
        name="分辨率",
        description="FFD控制点分辨率",
        items=[
            ('3x3x3', "3×3×3", "基础分辨率"),
            ('4x4x4', "4×4×4", "标准分辨率"), 
            ('5x5x5', "5×5×5", "高分辨率")
        ],
        default='4x4x4'
    )
    
    def execute(self, context):
        obj = context.active_object
        if not obj or obj.type != 'MESH':
            self.report({'ERROR'}, "请选择一个网格物体")
            return {'CANCELLED'}
        
        # 确保在物体模式
        if context.mode != 'OBJECT':
            bpy.ops.object.mode_set(mode='OBJECT')
        
        try:
            # 创建晶格
            bpy.ops.object.add(type='LATTICE')
            lattice = context.active_object
            lattice.name = f"FFD_Object_{self.resolution}"
            
            # 设置晶格分辨率
            res_map = {'3x3x3': (3,3,3), '4x4x4': (4,4,4), '5x5x5': (5,5,5)}
            res = res_map[self.resolution]
            lattice.data.points_u = res[0]
            lattice.data.points_v = res[1] 
            lattice.data.points_w = res[2]
            
            size ,center = get_bounds_center(obj)
            # 定位晶格到物体
            lattice.location = center
            scale = Vector((    max(size.x, min_size),    max(size.y, min_size),    max(size.z, min_size) ))
            
            lattice.scale = scale * 1.2
            
            # 为原始物体添加晶格修改器
            context.view_layer.objects.active = obj
            lattice_mod = obj.modifiers.new(name=f"FFD_Object_{self.resolution}", type='LATTICE')
            lattice_mod.object = lattice
            
            # 确保修改器在堆栈底部
            while obj.modifiers[0] != lattice_mod:
                bpy.ops.object.modifier_move_up(modifier=lattice_mod.name)
            
            self.report({'INFO'}, f"已为物体添加{self.resolution} FFD")
            bpy.context.view_layer.objects.active = lattice
            return {'FINISHED'}
            
        except Exception as e:
            self.report({'ERROR'}, f"添加物体FFD失败: {str(e)}")
            return {'CANCELLED'}

# ==================== 编辑模式FFD工具 ====================

class FFD_OT_add_edit_ffd(Operator):
    """为选中元素添加FFD"""
    bl_idname = "object.ffd_add_edit"
    bl_label = "选择FFD"
    bl_description = "为选中的顶点/边/面添加FFD修改器"
    bl_options = {'REGISTER', 'UNDO'}
    
    resolution: EnumProperty(
        name="分辨率",
        description="FFD控制点分辨率",
        items=[
            ('3x3x3', "3×3×3", "基础分辨率"),
            ('4x4x4', "4×4×4", "标准分辨率"), 
            ('5x5x5', "5×5×5", "高分辨率")
        ],
        default='4x4x4'
    )
    
    def get_selection_in_object_mode(self, obj):
        """在物体模式下获取选择数据"""
        selected_verts = []
        selected_faces = []
        
        # 切换到物体模式获取稳定数据
        bpy.ops.object.mode_set(mode='OBJECT')
        
        # 获取顶点选择
        for i, vert in enumerate(obj.data.vertices):
            if vert.select:
                selected_verts.append(i)
        
        # 获取面选择
        for i, face in enumerate(obj.data.polygons):
            if face.select:
                selected_faces.append(i)
        
        return selected_verts, selected_faces
    
    def execute(self, context):
        obj = context.active_object
        if not obj or obj.type != 'MESH':
            self.report({'ERROR'}, "请选择一个网格物体")
            return {'CANCELLED'}
        
        # 确保在物体模式开始
        if context.mode != 'OBJECT':
            bpy.ops.object.mode_set(mode='OBJECT')
        
        try:
            # 获取选择数据
            selected_verts, selected_faces = self.get_selection_in_object_mode(obj)
            
            if not selected_verts and not selected_faces:
                self.report({'ERROR'}, "请在编辑模式下选择一些顶点或面")
                return {'CANCELLED'}
            
            # 创建顶点组隔离选区
            vertex_group = obj.vertex_groups.new(name=f"FFD_Selection_{self.resolution}")
            
            # 添加选中的顶点
            if selected_verts:
                for vert_index in selected_verts:
                    vertex_group.add([vert_index], 1.0, 'REPLACE')
            
            # 添加选中的面
            if selected_faces:
                for face_index in selected_faces:
                    face = obj.data.polygons[face_index]
                    for vert_index in face.vertices:
                        vertex_group.add([vert_index], 1.0, 'REPLACE')
            
            # 计算选区边界
            world_verts = []
            for vert_index in selected_verts:
                vert = obj.data.vertices[vert_index]
                world_verts.append(obj.matrix_world @ vert.co)
            
            for face_index in selected_faces:
                face = obj.data.polygons[face_index]
                for vert_index in face.vertices:
                    vert = obj.data.vertices[vert_index]
                    world_verts.append(obj.matrix_world @ vert.co)
            
            # 移除重复顶点
            unique_verts = []
            for vert in world_verts:
                if vert not in unique_verts:
                    unique_verts.append(vert)
            
            if not unique_verts:
                self.report({'ERROR'}, "无法计算选择边界")
                return {'CANCELLED'}
            
            # 计算边界框
            min_co = Vector((
                min(v.x for v in unique_verts),
                min(v.y for v in unique_verts),
                min(v.z for v in unique_verts)
            ))
            max_co = Vector((
                max(v.x for v in unique_verts),
                max(v.y for v in unique_verts),
                max(v.z for v in unique_verts)
            ))
            
            center = (min_co + max_co) / 2
            size = max_co - min_co
            
            # 添加边距
            size = size * 1.2
            if size.length < 0.1:
                size = Vector((0.5, 0.5, 0.5))
            scale = Vector((    max(size.x, min_size),    max(size.y, min_size),    max(size.z, min_size) ))
                
            # 创建晶格
            bpy.ops.object.add(type='LATTICE')
            lattice = context.active_object
            lattice.name = f"FFD_Selection_{self.resolution}"
            
            # 设置晶格分辨率
            res_map = {'3x3x3': (3,3,3), '4x4x4': (4,4,4), '5x5x5': (5,5,5)}
            res = res_map[self.resolution]
            lattice.data.points_u = res[0]
            lattice.data.points_v = res[1] 
            lattice.data.points_w = res[2]
            
            # 定位晶格
            lattice.location = center
            lattice.scale = scale
            
            # 为原始物体添加晶格修改器
            context.view_layer.objects.active = obj
            lattice_mod = obj.modifiers.new(name=f"FFD_Selection_{self.resolution}", type='LATTICE')
            lattice_mod.object = lattice
            lattice_mod.vertex_group = vertex_group.name
            
            # 确保修改器在堆栈底部
            while obj.modifiers[0] != lattice_mod:
                bpy.ops.object.modifier_move_up(modifier=lattice_mod.name)
            
            context.view_layer.objects.active = lattice
            self.report({'INFO'}, f"已为选区添加{self.resolution} FFD")
            return {'FINISHED'}
            
        except Exception as e:
            self.report({'ERROR'}, f"添加选择FFD失败: {str(e)}")
            return {'CANCELLED'}

# ==================== 通用工具 ====================

class FFD_OT_collapse_ffd(Operator):
    """塌陷FFD修改器"""
    bl_idname = "object.ffd_collapse" 
    bl_label = "塌陷FFD"
    bl_description = "塌陷FFD修改器并保留其他修改器"
    bl_options = {'REGISTER', 'UNDO'}
    
    def execute(self, context):
        obj = context.active_object
        if not obj:
            self.report({'ERROR'}, "请选择一个物体")
            return {'CANCELLED'}
        
        # 确保在物体模式
        if context.mode != 'OBJECT':
            bpy.ops.object.mode_set(mode='OBJECT')
            
        try:
            # 查找FFD修改器
            ffd_mods = []
            for mod in obj.modifiers:
                if mod.type == 'LATTICE' and mod.name.startswith("FFD_"):
                    ffd_mods.append(mod)
                    
            if not ffd_mods:
                self.report({'ERROR'}, "未找到FFD修改器")
                return {'CANCELLED'}
            
            # 应用所有FFD修改器
            for ffd_mod in ffd_mods[:]:  # 使用副本遍历
                vertex_group_name = ffd_mod.vertex_group
                lattice_obj = ffd_mod.object
                
                # 应用修改器
                bpy.ops.object.modifier_apply(modifier=ffd_mod.name)
                
                # 清理顶点组
                if vertex_group_name:
                    vg = obj.vertex_groups.get(vertex_group_name)
                    if vg and vg.name.startswith("FFD_"):
                        obj.vertex_groups.remove(vg)
                
                # 清理晶格物体
                if lattice_obj and lattice_obj.name.startswith("FFD_"):
                    bpy.data.objects.remove(lattice_obj, do_unlink=True)
                
            self.report({'INFO'}, f"已塌陷 {len(ffd_mods)} 个FFD修改器")
            return {'FINISHED'}
            
        except Exception as e:
            self.report({'ERROR'}, f"塌陷FFD失败: {str(e)}")
            return {'CANCELLED'}

class FFD_OT_clear_all_ffd(Operator):
    """清理所有FFD数据"""
    bl_idname = "object.ffd_clear_all"
    bl_label = "清理FFD"
    bl_description = "清理所有FFD相关的修改器、顶点组和晶格"
    bl_options = {'REGISTER', 'UNDO'}
    
    def execute(self, context):
        obj = context.active_object
        if not obj:
            self.report({'ERROR'}, "请选择一个物体")
            return {'CANCELLED'}
        
        # 确保在物体模式
        if context.mode != 'OBJECT':
            bpy.ops.object.mode_set(mode='OBJECT')
            
        try:
            removed_count = 0
            
            # 移除FFD修改器
            ffd_mods = [mod for mod in obj.modifiers if mod.type == 'LATTICE' and mod.name.startswith("FFD_")]
            for mod in ffd_mods:
                obj.modifiers.remove(mod)
                removed_count += 1
            
            # 移除FFD顶点组
            ffd_vgs = [vg for vg in obj.vertex_groups if vg.name.startswith("FFD_")]
            for vg in ffd_vgs:
                obj.vertex_groups.remove(vg)
                removed_count += 1
            
            # 移除FFD晶格物体
            ffd_lattices = [obj for obj in bpy.data.objects if obj.type == 'LATTICE' and obj.name.startswith("FFD_")]
            for lattice in ffd_lattices:
                bpy.data.objects.remove(lattice, do_unlink=True)
                removed_count += 1
                
            self.report({'INFO'}, f"已清理 {removed_count} 个FFD项目")
            return {'FINISHED'}
            
        except Exception as e:
            self.report({'ERROR'}, f"清理失败: {str(e)}")
            return {'CANCELLED'}

# ==================== 面板界面 ====================

class FFD_PT_panel(Panel):
    """FFD工具面板"""
    bl_label = "FFD Tools"
    bl_idname = "FFD_PT_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "FFD Tools"
    
    def draw(self, context):
        layout = self.layout
        obj = context.active_object
        
        # 检查是否选中网格物体
        if not obj or obj.type != 'MESH':
            layout.label(text="请选择网格物体", icon='ERROR')
            return
        
        current_mode = context.mode
        
        # 物体模式工具
        if current_mode == 'OBJECT':
            box = layout.box()
            box.label(text="物体模式FFD", icon='OBJECT_DATA')
            
            row = box.row(align=True)
            row.operator("object.ffd_add_object", text="3×3×3").resolution = '3x3x3'
            row.operator("object.ffd_add_object", text="4×4×4").resolution = '4x4x4'
            row.operator("object.ffd_add_object", text="5×5×5").resolution = '5x5x5'
            
            box.label(text="为整个物体添加FFD", icon='INFO')
        
        # 编辑模式工具
        elif current_mode == 'EDIT_MESH':
            box = layout.box()
            box.label(text="编辑模式FFD", icon='EDITMODE_HLT')
            
            row = box.row(align=True)
            row.operator("object.ffd_add_edit", text="3×3×3").resolution = '3x3x3'
            row.operator("object.ffd_add_edit", text="4×4×4").resolution = '4x4x4'
            row.operator("object.ffd_add_edit", text="5×5×5").resolution = '5x5x5'
            
            box.label(text="为选中元素添加FFD", icon='INFO')
        
        # 通用工具
        box = layout.box()
        box.label(text="通用工具", icon='TOOL_SETTINGS')
        
        box.operator("object.ffd_collapse", text="塌陷所有FFD", icon='MOD_LATTICE')
        box.operator("object.ffd_clear_all", text="清理所有FFD", icon='TRASH')
        
        # 使用说明
        box = layout.box()
        box.label(text="使用说明", icon='QUESTION')
        box.label(text="物体模式: 为整个物体添加FFD")
        box.label(text="编辑模式: 选择点/面后添加FFD")
        box.label(text="塌陷: 保留对称和细分修改器")

# ==================== 注册 ====================

classes = (
    FFD_OT_add_object_ffd,
    FFD_OT_add_edit_ffd,
    FFD_OT_collapse_ffd,
    FFD_OT_clear_all_ffd,
    FFD_PT_panel,
)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)

if __name__ == "__main__":
    register()